from pydantic import BaseModel
import uvicorn
import fastapi
from fastapi.responses import StreamingResponse
import argparse
from models import ModelMiniCPMV4, ModelMiniCPMV4_5
import logging
import json

from logging_util import setup_root_logger

logger = logging.getLogger(__name__)


class Model:
    """模型管理类，负责加载和调用MiniCPM-V模型"""

    def __init__(self, model_path: str, model_type: str, instance_id: int = 0, gpu_id: int = None) -> None:
        """
        初始化模型实例

        Args:
            model_path: 模型文件路径或HuggingFace模型名称
            model_type: 模型类型 ('minicpmv4' 或 'minicpmv4_5')
            instance_id: 实例ID，用于多实例部署
            gpu_id: GPU设备ID，None表示使用默认GPU
        """
        self.instance_id = instance_id
        self.gpu_id = gpu_id

        # 设置GPU环境变量
        if gpu_id is not None:
            import os
            os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
            logger.info(f"实例 {instance_id}: 设置CUDA_VISIBLE_DEVICES={gpu_id}")

        logger.info(f"实例 {instance_id}: 初始化模型类型 {model_type}")

        # 根据模型类型选择对应的模型类
        match model_type.lower():
            case 'minicpmv4':
                self.model = ModelMiniCPMV4(model_path)
            case 'minicpmv4_5':
                self.model = ModelMiniCPMV4_5(model_path)
            case _:
                raise ValueError(f"不支持的模型类型: {model_type}")

        logger.info(f"实例 {instance_id}: 模型加载完成")

    def handler(self, query):
        """
        标准处理函数，同步返回模型推理结果

        Args:
            query: 包含图像、问题和参数的字典

        Returns:
            dict: 包含推理结果和token使用量的字典
        """
        # 调用模型进行推理
        res, output_tokens = self.model({
            "image": query["image"],           # 图像数据（base64编码）
            "question": query["question"],     # 用户问题
            "params": query.get("params", "{}"),  # 模型参数（JSON字符串）
            "temporal_ids": query.get("temporal_ids", None)  # 时序ID（用于视频处理）
        })
        return {
            "result": res,                    # 模型推理结果
            "usage": {"output_tokens": output_tokens}  # 输出token数量统计
        }

    def stream_handler(self, query):
        """
        流式处理函数，返回生成器用于实时输出

        Args:
            query: 包含图像、问题和参数的字典

        Returns:
            generator: 模型输出的生成器
        """
        # 解析并修改参数，启用流式输出
        params = json.loads(query.get("params", "{}"))
        params["stream"] = True  # 启用流式传输
        query["params"] = json.dumps(params)

        # 调用模型获取流式生成器
        generator = self.model({
            "image": query["image"],
            "question": query["question"],
            "params": query["params"],
            "temporal_ids": query.get("temporal_ids", None)
        })

        return generator


class Item(BaseModel):
    """API请求数据模型定义"""
    image: str           # 图像数据（base64编码字符串）
    question: str        # 用户提问的问题
    params: str         # 模型参数（JSON字符串）
    temporal_ids: str = None  # 时序ID（可选，用于视频帧处理）

# 全局变量
model:Model = None  # 模型实例
args = None        # 命令行参数

def initialize_server():
    """初始化服务器配置和模型加载"""
    global model, args

    # 命令行参数解析
    parser = argparse.ArgumentParser(description='Server for MiniCPM-V')
    parser.add_argument('--port', type=int, default=9999,
                        help='Port to run the server on')
    parser.add_argument('--log_dir', type=str, default='logs',
                        help='Directory for log files')
    parser.add_argument('--model_path', type=str, default='openbmb/MiniCPM-V-4_5',
                        help='Path to the model directory')
    parser.add_argument('--model_type', type=str, default='minicpmv4_5',
                        help='Type of the model to use')
    parser.add_argument('--instance_id', type=int, default=0,
                        help='Instance ID for multi-instance deployment')
    parser.add_argument('--gpu_id', type=int, default=None,
                        help='GPU ID to use for this instance')
    args = parser.parse_args()

    # 配置日志系统
    setup_root_logger(local_dir=args.log_dir)

    # 打印启动信息
    logger.info(f"="*50)
    logger.info(f"启动MiniCPM-V服务实例")
    logger.info(f"实例ID: {args.instance_id}")
    logger.info(f"GPU ID: {args.gpu_id}")
    logger.info(f"端口: {args.port}")
    logger.info(f"模型路径: {args.model_path}")
    logger.info(f"模型类型: {args.model_type}")
    logger.info(f"日志目录: {args.log_dir}")
    logger.info(f"="*50)

    # 初始化模型实例
    model = Model(args.model_path, args.model_type, args.instance_id, args.gpu_id)

# 创建FastAPI应用
app = fastapi.FastAPI()


@app.get("/")
def read_root():
    """根路径接口，返回服务器状态信息"""
    return {
        "message": "MiniCPM-V server",
        "instance_id": args.instance_id,
        "gpu_id": args.gpu_id,
        "port": args.port,
        "model_type": args.model_type,
        "status": "running"
    }


@app.post("/api")
def websocket(item: Item):
    """标准API接口，同步处理请求并返回完整结果"""
    logger.info(f'params: {str(item.params)}')  # 记录请求参数
    query = item.dict()  # 转换为字典格式
    res = model.handler(query)  # 调用模型处理

    logger.info(f'result: {str(res)}')  # 记录返回结果
    return {'data': res}  # 返回JSON格式结果


@app.post("/api/stream")
def stream_api(item: Item):
    """流式API接口，实时返回模型生成结果"""
    query = item.dict()

    def event_generator():
        """事件生成器，用于流式输出"""
        try:
            # 获取模型流式生成器
            generator = model.stream_handler(query)
            full_response = ""  # 完整响应内容
            output_tokens = 0   # token计数器

            # 遍历生成器的每个输出块
            for chunk in generator:
                full_response += chunk
                output_tokens += 1
                # 构建流式数据包
                data = {
                    "chunk": chunk,           # 当前输出块
                    "full_response": full_response,  # 累计完整响应
                    "finished": False         # 是否结束标志
                }
                # 以Server-Sent Events格式发送数据
                yield f"data: {json.dumps(data, ensure_ascii=False)}\n\n"

            # 发送结束信号
            final_data = {
                "chunk": "",
                "full_response": full_response,
                "finished": True
            }
            yield f"data: {json.dumps(final_data, ensure_ascii=False)}\n\n"

        except Exception as e:
            # 错误处理
            error_data = {
                "error": str(e),
                "finished": True
            }
            logger.error(f"Stream error: {e}")
            yield f"data: {json.dumps(error_data, ensure_ascii=False)}\n\n"

    # 返回流式响应
    return StreamingResponse(
        event_generator(),
        media_type="text/plain",  # 媒体类型
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Access-Control-Allow-Origin": "*",  # CORS配置
            "Access-Control-Allow-Headers": "*",
        }
    )

if __name__ == "__main__":
    # 主程序入口
    initialize_server()  # 初始化服务器

    # 配置并启动uvicorn服务器
    _cfg = uvicorn.Config(app, host="0.0.0.0", port=args.port, workers=2)
    uvicorn.Server(_cfg).run()